import numpy as np
import math
from utils.grid_map import GridMap

class DWAPlanner:
    def __init__(self, max_speed=1.5, max_accel=0.5, max_yawrate=40.0*math.pi/180.0, 
                 max_dyawrate=40.0*math.pi/180.0, v_reso=0.1, yawrate_reso=0.1*math.pi/180.0, 
                 dt=0.1, predict_time=1.0, to_goal_cost_gain=1.0, speed_cost_gain=0.1, 
                 obstacle_cost_gain=1.0, robot_radius=1.0):
        """
        DWA (Dynamic Window Approach) local planner
        :param max_speed: maximum speed (m/s)
        :param max_accel: maximum acceleration (m/s²)
        :param max_yawrate: maximum yaw rate (rad/s)
        :param max_dyawrate: maximum yaw rate change (rad/s²)
        :param v_reso: speed resolution (m/s)
        :param yawrate_reso: yaw rate resolution (rad/s)
        :param dt: time step (s)
        :param predict_time: prediction time (s)
        :param to_goal_cost_gain: goal cost gain
        :param speed_cost_gain: speed cost gain
        :param obstacle_cost_gain: obstacle cost gain
        :param robot_radius: robot radius (m)
        """
        self.max_speed = max_speed
        self.max_accel = max_accel
        self.max_yawrate = max_yawrate
        self.max_dyawrate = max_dyawrate
        self.v_reso = v_reso
        self.yawrate_reso = yawrate_reso
        self.dt = dt
        self.predict_time = predict_time
        self.to_goal_cost_gain = to_goal_cost_gain
        self.speed_cost_gain = speed_cost_gain
        self.obstacle_cost_gain = obstacle_cost_gain
        self.robot_radius = robot_radius

    def plan(self, current_state, goal, grid_map, dynamic_obstacles=[]):
        """
        DWA path planning
        :param current_state: [x, y, yaw, v, yawrate]
        :param goal: [x, y] in world coordinates
        :param grid_map: GridMap object
        :param dynamic_obstacles: list of dynamic obstacle positions [(x,y), ...]
        :return: best velocity (v, yawrate)
        """
        x, y, yaw, v, yawrate = current_state
        dw = self.calc_dynamic_window(x, y, yaw, v, yawrate)
        best_u, best_trajectory = self.search_best_trajectory(dw, current_state, goal, grid_map, dynamic_obstacles)
        return best_u, best_trajectory

    def calc_dynamic_window(self, x, y, yaw, v, yawrate):
        """Calculate dynamic window based on current state and constraints"""
        # Dynamic window from motion constraints
        Vs = [0, self.max_speed, -self.max_yawrate, self.max_yawrate]

        # Dynamic window from acceleration constraints
        Vd = [
            v - self.max_accel * self.dt,
            v + self.max_accel * self.dt,
            yawrate - self.max_dyawrate * self.dt,
            yawrate + self.max_dyawrate * self.dt
        ]

        # Merge dynamic windows
        dw = [
            max(Vs[0], Vd[0]),
            min(Vs[1], Vd[1]),
            max(Vs[2], Vd[2]),
            min(Vs[3], Vd[3])
        ]

        return dw

    def predict_trajectory(self, x, y, yaw, v, yawrate, dt, predict_time):
        """Predict trajectory with current velocity"""
        trajectory = []
        for _ in range(int(predict_time/dt)):
            x += v * math.cos(yaw) * dt
            y += v * math.sin(yaw) * dt
            yaw += yawrate * dt
            v = max(0, min(v, self.max_speed))  # Ensure speed doesn't exceed limits
            trajectory.append([x, y, yaw, v])
        return trajectory

    def evaluate_trajectory(self, trajectory, goal, grid_map, dynamic_obstacles):
        """Evaluate trajectory cost"""
        # Goal cost
        goal_cost = self.to_goal_cost_gain * self.calc_goal_cost(trajectory, goal)

        # Speed cost (higher speed is better)
        speed_cost = self.speed_cost_gain * (self.max_speed - trajectory[-1][3])

        # Obstacle cost
        obstacle_cost = self.obstacle_cost_gain * self.calc_obstacle_cost(trajectory, grid_map, dynamic_obstacles)

        # Total cost
        total_cost = goal_cost + speed_cost + obstacle_cost
        return total_cost, goal_cost, speed_cost, obstacle_cost

    def calc_goal_cost(self, trajectory, goal):
        """Calculate cost to goal"""
        dx = goal[0] - trajectory[-1][0]
        dy = goal[1] - trajectory[-1][1]
        return math.hypot(dx, dy)

    def calc_obstacle_cost(self, trajectory, grid_map, dynamic_obstacles):
        """Calculate cost to obstacles"""
        min_dist = float('inf')
        robot_radius = self.robot_radius

        # Check static obstacles
        for (x, y, _, _) in trajectory:
            grid_pos = grid_map.world2grid([x, y])
            if grid_map.is_occupied(grid_pos[0], grid_pos[1]):
                return float('inf')  # Collision

            # Check distance to nearby obstacles
            for i in range(-2, 3):
                for j in range(-2, 3):
                    nx = grid_pos[0] + i
                    ny = grid_pos[1] + j
                    if grid_map.is_occupied(nx, ny):
                        dist = math.hypot(nx*grid_map.resolution - x, ny*grid_map.resolution - y)
                        if dist < min_dist:
                            min_dist = dist

        # Check dynamic obstacles
        for (ox, oy) in dynamic_obstacles:
            for (x, y, _, _) in trajectory:
                dist = math.hypot(ox - x, oy - y)
                if dist < robot_radius:
                    return float('inf')  # Collision with dynamic obstacle
                if dist < min_dist:
                    min_dist = dist

        if min_dist == float('inf'):  # No obstacles nearby
            return 0
        return 1.0 / min_dist  # Higher cost for closer obstacles

    def search_best_trajectory(self, dw, current_state, goal, grid_map, dynamic_obstacles):
        """Search best trajectory in dynamic window"""
        x, y, yaw, v, yawrate = current_state
        best_cost = float('inf')
        best_u = [0.0, 0.0]
        best_trajectory = []

        # Sample velocity space
        for v_sample in np.arange(dw[0], dw[1], self.v_reso):
            for yawrate_sample in np.arange(dw[2], dw[3], self.yawrate_reso):
                trajectory = self.predict_trajectory(x, y, yaw, v_sample, yawrate_sample, self.dt, self.predict_time)
                cost, goal_cost, speed_cost, obstacle_cost = self.evaluate_trajectory(trajectory, goal, grid_map, dynamic_obstacles)

                # Update best trajectory
                if cost < best_cost:
                    best_cost = cost
                    best_u = [v_sample, yawrate_sample]
                    best_trajectory = trajectory

        return best_u, best_trajectory